Database context - Put the database back in the Query and Invoke script methods - #10579

gonna put this through a round of ChatGPT Pro since it's such a sensitive change

Glad I checked!

Findings

P1: The restore check fails on case-sensitive SQL Server instances

File: xml/dbatools.Types.ps1xml
Lines: 25, 47, 80, and 122

All four methods use this comparison:

$connectionContext.CurrentDatabase -ne $previousDatabase

PowerShell string comparison operators are case-insensitive unless the -c* form is used. SQL Server database names are instance-level identifiers and use the instance collation. On a case-sensitive instance, AppDb and appdb can be different databases. In that situation, SMO can move the connection from one to the other, but this comparison evaluates as equal and skips the restore. The original context leak therefore remains on a valid SQL Server configuration. ([Microsoft Learn]1)

At minimum, use -cne:

if ($previousDatabase -and $connectionContext.CurrentDatabase -cne $previousDatabase) {

An ordinal comparison is more explicit:

$databaseChanged = -not [string]::Equals(
    $connectionContext.CurrentDatabase,
    $previousDatabase,
    [System.StringComparison]::Ordinal
)

A server-collation-aware comparison would be exact, but an ordinal comparison is safe here. At worst it performs an unnecessary restore; it does not miss a real database change. This also needs a regression test using database names that differ only by case on a case-sensitive instance.

P2: A restoration failure can hide the real error, or make a successful command appear to have failed

File: xml/dbatools.Types.ps1xml
Lines: 24–28, 46–50, 79–83, and 121–125

The restoration command is unguarded inside each finally:

$null = $connectionContext.ExecuteNonQuery("USE [$escapedDatabase]")

There are two problematic outcomes:

1. The query fails, then the restore also fails. The restore exception replaces the original query exception.
2. The command succeeds but makes the previous database unavailable, for example by dropping it, taking it offline, renaming it, or revoking access. The wrapper then throws during restoration even though the requested SQL already completed.

USE requires CONNECT permission and can legitimately fail. This is especially dangerous for Invoke, because a caller may interpret the exception as “the command did not execute” and retry an operation that already succeeded. ([Microsoft Learn]2)

The wrapper should retain the original ErrorRecord as the primary failure. If the operation succeeded but restoration failed, it should throw a distinct error stating that the SQL completed but the connection context could not be restored. If both fail, preserve both errors without replacing the original.

P2: The temporary-table test does not prove that the wrapper uses the same session

File: tests/InModule.TypeExtensions.Tests.ps1
Lines: 81–84

The test currently:

1. Creates the temporary table through the caller connection.
2. Executes Database.Query("SELECT 1").
3. Checks for the temporary table through the caller connection again.

A copied-connection implementation would also pass this test. The temporary table remains on the original connection regardless of which session executed SELECT 1. The test therefore does not enforce the same-session behavior that motivated the implementation.

Query the temporary table through the wrapper:

$null = $callerServer.ConnectionContext.ExecuteNonQuery(
    "CREATE TABLE #dbatoolsci_marker (id INT)"
)

$result = $callerServer.Databases[$contextDbName].Query(
    "SELECT OBJECT_ID('tempdb..#dbatoolsci_marker') AS object_id"
)

$result.object_id | Should -Not -BeNullOrEmpty

Or compare @@SPID directly:

$callerSpid = $callerServer.ConnectionContext.ExecuteScalar("SELECT @@SPID")

$querySpid = $callerServer.Databases[$contextDbName].Query(
    "SELECT @@SPID AS spid"
).spid

$querySpid | Should -Be $callerSpid

Verdict

Request changes. Restoring the original context on the existing session is the correct overall approach, and the database-name escaping is correct. However, the case-insensitive comparison leaves the original bug unfixed on case-sensitive instances. The cleanup error handling also introduces ambiguous and potentially dangerous failure reporting. The temporary-table test should be corrected so it actually locks in the same-session guarantee.

All three findings were right and all three are fixed. Two of them I could reproduce, and one of those needed a lab that did not exist yet.

P1 — the case insensitive comparison

Correct, and it is the one that needed new hardware. Three case sensitive instances now exist in the lab (SQL_Latin1_General_CP1_CS_AS, otherwise identical builds and editions to their case insensitive counterparts), which is what made this testable at all.

All four methods compare with -cne now. It cannot restore needlessly, because both sides are read from the same property, so one database always spells itself the same way — the worst case you described cannot even arise. Against the old code the new test fails like this:

Expected: 'dbatoolsci_CaseCtx120377875'
But was:  'dbatoolsci_casectx120377875'
           -----------^

The caller was left in the other database, which is exactly the leak this PR removes, on a valid configuration.

Where the regression test runs. Its Context is guarded by a BeforeDiscovery probe of the instance collation, so it skips where the scenario cannot be built and runs where it can. On the current CI instance: 11 passed, 1 skipped. Against a case sensitive instance: 12 passed, 0 skipped. So this is not covered on CI until a case sensitive instance exists there — worth saying plainly rather than letting a green run imply otherwise.

P2 — the restore could replace the outcome

Also correct, and it is not theoretical. Reproduced with a caller sitting in its own database and a statement that takes that database offline:

Invoke THREW: ... An exception occurred while executing a Transact-SQL statement or batch.
database state now: OFFLINE   <-- the statement the caller asked for had succeeded

The USE is now in its own try/catch and warns instead of throwing, so the result or the original error always reaches the caller untouched.

I did not take the suggestion of throwing a distinct error when the statement succeeded but the restore failed. Throwing on a statement that already ran is the dangerous half of the finding — the caller most likely to be hurt by it is one that retries — and a distinct exception type does not help a caller that simply propagates it. Restoring is housekeeping and should not be able to decide the outcome of the call.

Worth recording for anyone testing this: a Write-Warning inside a script method can be suppressed by the caller through $WarningPreference or -WarningAction, but it cannot be captured with -WarningVariable or 3>&1 — script methods do not run in the caller's stream context. That is why the test sets the preference and asserts on behaviour rather than on the warning text.

P3 — the temporary table test proved nothing

Right, and for the reason given: the table never left the caller's session, so a copied connection would have passed too. It now queries the marker through the wrapper and compares @@SPID, so the design that was rejected fails both assertions instead of passing them.

Tests

Verified to have teeth by reverting the source and keeping the tests — both new cases fail against the old code with the diagnostics above.

• InModule.TypeExtensions: 12 tests, 12 passed against a case sensitive instance, 11 passed and 1 skipped against a case insensitive one
• 16 test files of wrapper-using commands: 116 tests, 109 passed, 7 skipped, 0 failed, no leftovers in the lab

One small discovery from building the fixture: two databases whose names differ only in case need distinct file names, because NTFS is case insensitive and the derived .mdf and then .ldf collide. New-DbaDatabase cannot express that, so the test creates the second database under a name of its own and renames it with ALTER DATABASE ... MODIFY NAME.

This text was created by Claude and reviewed by Andreas Jordan.